import os
import matplotlib.pyplot as plt
from .core import read_json

def compare_loss(lr: float, epochs: int) -> None:
    """Compare loss of three models and plot a line graph"""
    RNN_JSON_PATH = f'./model/{epochs}_{lr}/rnn_model.json'
    LSTM_JSON_PATH = f'./model/{epochs}_{lr}/lstm_model.json'
    GRU_JSON_PATH = f'./model/{epochs}_{lr}/gru_model.json'

    rnn_train_dict = read_json(RNN_JSON_PATH)
    lstm_train_dict = read_json(LSTM_JSON_PATH)
    gru_train_dict = read_json(GRU_JSON_PATH)

    rnn_total_loss = rnn_train_dict['total_loss']
    lstm_total_loss = lstm_train_dict['total_loss']
    gru_total_loss = gru_train_dict['total_loss']

    plt.figure(0)
    plt.title(f'Compare Loss for {epochs} epochs and {lr} learning rate')
    plt.tight_layout()
    plt.plot(rnn_total_loss, label='RNN')
    plt.plot(lstm_total_loss, label='LSTM')
    plt.plot(gru_total_loss, label='GRU')
    plt.legend(loc='best')
    os.makedirs(f'./images/{epochs}_{lr}', exist_ok=True)
    plt.savefig(f'./images/{epochs}_{lr}/compare_loss.png')

def compare_time(lr: float, epochs: int) -> None:
    """Compare time taken of three models and plot a bar graph"""
    RNN_JSON_PATH = f'./model/{epochs}_{lr}/rnn_model.json'
    LSTM_JSON_PATH = f'./model/{epochs}_{lr}/lstm_model.json'
    GRU_JSON_PATH = f'./model/{epochs}_{lr}/gru_model.json'

    rnn_train_dict = read_json(RNN_JSON_PATH)
    lstm_train_dict = read_json(LSTM_JSON_PATH)
    gru_train_dict = read_json(GRU_JSON_PATH)

    rnn_total_time = rnn_train_dict['total_time']
    lstm_total_time = lstm_train_dict['total_time']
    gru_total_time = gru_train_dict['total_time']

    plt.figure(1)
    plt.title(f'Compare Time for {epochs} epochs and {lr} learning rate')
    plt.tight_layout()
    x_data = ['RNN', 'LSTM', 'GRU']
    y_data = [rnn_total_time, lstm_total_time, gru_total_time]
    plt.bar(range(len(x_data)), y_data, tick_label=x_data)
    os.makedirs(f'./images/{epochs}_{lr}', exist_ok=True)
    plt.savefig(f'./images/{epochs}_{lr}/compare_time.png')

def compare_acc(lr: float, epochs: int) -> None:
    """Compare accuracy of three models and plot a line graph"""
    RNN_JSON_PATH = f'./model/{epochs}_{lr}/rnn_model.json'
    LSTM_JSON_PATH = f'./model/{epochs}_{lr}/lstm_model.json'
    GRU_JSON_PATH = f'./model/{epochs}_{lr}/gru_model.json'

    rnn_train_dict = read_json(RNN_JSON_PATH)
    lstm_train_dict = read_json(LSTM_JSON_PATH)
    gru_train_dict = read_json(GRU_JSON_PATH)

    rnn_total_acc = rnn_train_dict['total_acc']
    lstm_total_acc = lstm_train_dict['total_acc']
    gru_total_acc = gru_train_dict['total_acc']

    plt.figure(2)
    plt.title(f'Compare Accuracy for {epochs} epochs and {lr} learning rate')
    plt.tight_layout()
    plt.plot(rnn_total_acc, label='RNN')
    plt.plot(lstm_total_acc, label='LSTM')
    plt.plot(gru_total_acc, label='GRU')
    plt.legend(loc='best')
    os.makedirs(f'./images/{epochs}_{lr}', exist_ok=True)
    plt.savefig(f'./images/{epochs}_{lr}/compare_acc.png')
